from jqdatasdk.api import get_bars
import pandas as pd
from datetime import datetime,date,time,timedelta
from util.query_data.control.ck_control import *

class ProcessedQueryData():
    
    #初始化
    #source_type: jk 聚宽 jj 掘金 mk 米筐
    def __init__(self,source_type:str):
        self.ck = ClickHouse()
        self.source_type = source_type
        self.source_database_name = "processed_data_stocks"
        self.trade_date_table_name = "stocks_trade_date_"
        self.code_info_table_name = "stocks_code_info_"
        self.md_bar_day_table_name = "stocks_md_bar_day_"
        self.financial_market_capitalization_table_name = "stocks_financial_market_capitalization_"
        self.mtss_info_table_name = "stocks_mtss_info_"
        self.stk_hold_info_table_name = "stocks_stk_hold_info_"
        self.profile_xr_xd_table_name = "stocks_profile_xr_xd_"

    def read_data(self,que_dict:dict,table_name:str):
        que_table_name = table_name+self.source_type
        result_data = self.ck.read_dataframe(self.source_database_name,que_table_name,que_dict)
        return result_data 

    """
    get_trade_date
    获取交易日数据

    start_date:开始时间
    end_date:结束时间

    返回 交易日dataframe
    """
    def get_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        trade_dataframe = self.read_data(que_dict,self.trade_date_table_name)
        return trade_dataframe
    
    """
    get_last_n_trade_date
    获取前几个交易日

    day_date:当前日期
    day_num:几个交易日

    返回 交易日
    """
    def get_last_n_trade_date(self,day_date:str,day_num:int):
        que_dict = {f"trade_date":"<'"+str(day_date)+"'"}
        trade_dataframe = self.read_data(que_dict,self.trade_date_table_name)
        trade_dataframe = trade_dataframe.sort_values(by=['trade_date'],ascending=False)
        return trade_dataframe.iloc[day_num-1].trade_date

    """
    get_next_n_trade_date
    获取后几个交易日

    day_date:当前日期
    day_num:几个交易日

    返回 交易日
    """
    def get_next_n_trade_date(self,day_date:str,day_num:int):
        que_dict = {f"trade_date":">'"+str(day_date)+"'"}
        trade_dataframe = self.read_data(que_dict,self.trade_date_table_name)
        trade_dataframe = trade_dataframe.sort_values(by=['trade_date'],ascending=True)
        return trade_dataframe.iloc[day_num-1].trade_date
    
    """
    get_code_info_by_trade_date

    获取股票代码-根据交易日

    start_date:开始时间
    end_date:结束时间

    返回 股票代码
    """
    def get_code_info_by_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        result_dataframe = self.read_data(que_dict,self.code_info_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        return result_dataframe
    
    """
    get_md_bar_day_by_trade_date

    获取股票行情日k-根据交易日

    start_date:开始时间
    end_date:结束时间

    返回 股票代码
    """
    def get_md_bar_day_by_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        result_dataframe = self.read_data(que_dict,self.md_bar_day_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        return result_dataframe

    """
    get_md_bar_day_xr_xd_by_trade_date

    获取股票行情日k-后复权-根据交易日

    start_date:开始时间
    end_date:结束时间

    返回 股票代码
    """
    def get_md_bar_day_xr_xd_by_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        result_dataframe = self.read_data(que_dict,self.md_bar_day_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        xr_xd_list_a = self.get_profile_xr_xd_by_trade_date(start_date,end_date)
        select_sec_data = pd.DataFrame(columns=result_dataframe.columns)
        groupby_code = result_dataframe.groupby("code")
        for key in groupby_code.groups:
            sec_data = groupby_code.get_group(key)
            sec_data_xd = self.get_xr_xd_dataframe(sec_data,key,xr_xd_list_a)
            select_sec_data = select_sec_data.append(sec_data_xd)
        return select_sec_data

    def get_xr_xd_dataframe(self,sec_data_day_a,code,xr_xd_list_a):
        sec_data_day_a = sec_data_day_a.sort_values(by='trade_date',ascending=True)
        str_date_f_q = sec_data_day_a.iloc[0]
        end_date_f_q = sec_data_day_a.iloc[len(sec_data_day_a)-1]
        xr_xd_list = xr_xd_list_a[(xr_xd_list_a['trade_date']>=str_date_f_q.trade_date)&(xr_xd_list_a['trade_date']<=end_date_f_q.trade_date)&(xr_xd_list_a['code']==code)]
        if len(xr_xd_list)>0:
            for index, row in xr_xd_list.iterrows():
                try:
                    xr_xd = row
                    # if np.isnan(xr_xd.float_capital_after_transfer)==True:
                    #     continue
                    xr_xd_yz = 1+xr_xd.share_trans_ratio+xr_xd.share_div_ratio
                    new_list = pd.DataFrame(columns=sec_data_day_a.columns)
                    new_list = new_list.append(sec_data_day_a[sec_data_day_a['trade_date']<xr_xd.trade_date])
                    fq_list = sec_data_day_a[sec_data_day_a['trade_date']>=xr_xd.trade_date]
                    fq_list['fh'] = round(xr_xd.cash_div,3)
                    fq_list['xr_xd_yz'] = xr_xd_yz
                    fq_list['open'] = round(fq_list['open']*fq_list['xr_xd_yz']+fq_list['fh'],2)
                    fq_list['close'] = round(fq_list['close']*fq_list['xr_xd_yz']+fq_list['fh'],2)
                    fq_list['high'] = round(fq_list['high']*fq_list['xr_xd_yz']+fq_list['fh'],2)
                    fq_list['low'] = round(fq_list['low']*fq_list['xr_xd_yz']+fq_list['fh'],2)
                    new_list = new_list.append(fq_list.drop(columns=['fh','xr_xd_yz']))
                    sec_data_day_a = new_list
                except Exception as e:
                    print("出错了")
                    continue
            return sec_data_day_a
        else:
            return sec_data_day_a


    """
    get_financial_market_capitalization_by_trade_date

    获取股票市值-根据交易日

    start_date:开始时间
    end_date:结束时间

    返回 股票市值
    """
    def get_financial_market_capitalization_by_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        result_dataframe = self.read_data(que_dict,self.financial_market_capitalization_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        return result_dataframe

    """
    get_mtss_info_by_trade_date

    获取融资融券-根据交易日

    start_date:开始时间
    end_date:结束时间

    返回 融资融券
    """
    def get_mtss_info_by_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        result_dataframe = self.read_data(que_dict,self.mtss_info_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        return result_dataframe

    """
    get_stk_hold_info_by_trade_date

    获取北向资金持仓-根据交易日

    start_date:开始时间
    end_date:结束时间
    type:SZ 深股通 SH 沪股通 XG 港股通

    返回 北向资金持仓
    """
    def get_stk_hold_info_by_trade_date(self,start_date:str,end_date:str,type:Union[str, None] = None):
        if type==None:
            que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        else:
            que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' and link_id='"+type+"' "}
        result_dataframe = self.read_data(que_dict,self.stk_hold_info_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        return result_dataframe

    """
    get_profile_xr_xd_by_trade_date

    获取除权除息-根据交易日

    start_date:开始时间
    end_date:结束时间

    返回 除权除息
    """
    def get_profile_xr_xd_by_trade_date(self,start_date:str,end_date:str):
        que_dict = {f"trade_date":">='"+str(start_date)+"' and trade_date <='"+str(end_date)+"' "}
        result_dataframe = self.read_data(que_dict,self.profile_xr_xd_table_name)
        result_dataframe = result_dataframe.sort_values(by=['trade_date'],ascending=True)
        return result_dataframe